from convdiffsolverbase import ConvDiffSolverBase
from dolfin import *
import numpy

class Codina(ConvDiffSolverBase):
    '''Nonlinear crosswind diffusion method by Codina(1993) - A discontinuity-capturing crosswind dissipation... .'''
    def __init__(self):
        self.name = 'Codina'
        ConvDiffSolverBase.__init__(self)

    def solve(self,problem,get_tau,get_sigma):
        '''Function get_sigma for linear case to obtain initial guess.'''
        
        problemName, mesh, eps, BCValues, b, f, BCIndicators = problem.get_vars()
        
        V = FunctionSpace(mesh,'CG',1) 
        VV = VectorFunctionSpace(mesh,'CG',1)
        DGV = FunctionSpace(mesh,'DG',0)
        
        u = TrialFunction(V)
        u_ = Function(V)
        varphi = TestFunction(V)

        v = interpolate(b,VV)
        bcs = [DirichletBC(V,value,where) for value, where in zip(BCValues, BCIndicators)]

        #---------------------------------------------------------------------------------------------------
        def get_sigma_linear(v,eps,sigma_): # for u_ 
            sigma_.vector()[:] = get_sigma(v,eps)
        
        def get_sigma_nonlinear(f,v,u_,eps,sigma_): # for nonlinear case
            phi = TestFunction(DGV)  # universal test function   
           
            solSem = TrialFunction(DGV) # trial function for seminorm and norm of solution and residuum 
            solNorm = TrialFunction(DGV) 
            res = TrialFunction(DGV)
        
            solSemF = Function(DGV) #this is where the values will go
            resF = Function(DGV)
            solNormF = Function(DGV)            
        
            solSemA = solSem*phi*dx
            solSemb = sqrt(inner(nabla_grad(u_),nabla_grad(u_)))*phi*dx 
            solve(solSemA==solSemb,solSemF)   # W1,2 seminorm of u_
            
            resA = res*phi*dx
            resb = (-eps*nabla_div(nabla_grad(u_))+inner(v,nabla_grad(u_))-f)*phi*dx
            solve(resA==resb,resF)            # residuum of u_

            solNormA = solNorm*phi*dx
            solNormb = sqrt(resF*resF)*phi*dx   
            solve(solNormA==solNormb,solNormF)# W0,2(L2) norm of residuum of u_
            
            nc = mesh.num_cells()
            sigmaValues = numpy.zeros(nc)
            beta = 0.7
            for i in xrange(nc): # build the values
                c = Cell(mesh,i)
               
                diamK = c.diameter()  #elements diameter
                x = c.midpoint()
                
                norm = solNormF(x.x(),x.y())
                seminorm = solSemF(x.x(),x.y())
                
                if norm < 1E-16 or seminorm < 1E-16:
                    sigmaValues[i] = 0.
                else:
                    sigmaValues[i] = 0.5*max(0,beta - 2*eps*seminorm/norm/diamK)*diamK*norm/seminorm
                
            sigma_.vector()[:] = sigmaValues
        #-------------------------------------------------------------------------------------------------- 

        tau = Function(DGV)
        tau.vector()[:] = get_tau(v,eps) # should be the same as in upwind diffusion
        
        sigma_ = Function(DGV)
       
        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                          ) # a way to realize swith on |v| == 0.
                        
        pMatrix = conditional(
                             gt(sqrt(inner(v,v)),1E-16),
                             Identity(2)-outer(v,v)/inner(v,v),
                             Constant(0.)*Identity(2),
                             ) # I - v X v
        
        a = eps*inner(nabla_grad(u),nabla_grad(varphi))*dx + inner(v,nabla_grad(u))*varphi*dx +\
            tau*pert*inner(dot(nabla_grad(u),outer(v,v)),nabla_grad(varphi))*dx+\
            sigma_*inner(dot(nabla_grad(u),pMatrix),nabla_grad(varphi))*dx       

        l = f*varphi*dx + tau*pert*f*inner(v,nabla_grad(varphi))*dx           # universal weak form
 
        get_sigma_linear(v,eps,sigma_)
        solve(a==l,u_,bcs)                   # u_ holds solution to linear problem, initial guess
        
        iter = 0
        maxIter = 50
        error = 1
        tol = 1E-3
        u = Function(V)
        
        while error > tol and iter<maxIter:
            iter += 1
            
            get_sigma_nonlinear(f,v,u_,eps,sigma_)
            A,L = assemble_system(a,l,bcs)
            solve(A,u.vector(),L)
            
            diff = u.vector().array() - u_.vector().array()
            error = numpy.linalg.norm(diff, ord=numpy.Inf)
            u_.vector()[:] = u_.vector().array() + 1.0*diff
            
            print iter, error
        
        plot(u_,interactive=True)
        self.save(u,problem)
